package org.mmtk.plan.scale.balance;

import org.mmtk.plan.scale.utils.ArayHelper;
import org.mmtk.utility.Log;
import org.mmtk.utility.options.Options;
import org.mmtk.vm.VM;
import org.vmmagic.pragma.Inline;
import org.vmmagic.pragma.Uninterruptible;
import org.vmmagic.unboxed.ObjectReference;
import org.vmmagic.unboxed.ObjectReferenceArray;

@Uninterruptible
public class TraceBranchBalancer implements BalancerInterface {
	private static final boolean DEBUG = false;
	/* constants */
	private static final int HISTORY_SIZE = 50000;

	/* structures */
	private ObjectReferenceArray objectRefs;
	private ObjectReferenceArray candidateRefs;
	private long[] candidateDepth;
	private long[] level;
	private int[] children;
	private long[] descendants;
	private long[] toBottom;

	/* counters */
	private int size;
	private long lost;
	private long setPointers;
	private long maxDepth;

	TraceBranchBalancer() {
		this.objectRefs = ObjectReferenceArray.create(HISTORY_SIZE);
		this.candidateRefs = ObjectReferenceArray.create(HISTORY_SIZE);
		this.candidateDepth = new long[HISTORY_SIZE];
		this.level = new long[HISTORY_SIZE];
		this.children = new int[HISTORY_SIZE];
		this.descendants = new long[HISTORY_SIZE];
		this.toBottom = new long[HISTORY_SIZE];
		initBalance();
	}

	@Inline
	public void initBalance() {
		this.size = 0;
		this.lost = 0;
		this.setPointers = 0;
		this.maxDepth = 0;
		ArayHelper.zeroArray(this.level);
		ArayHelper.zeroArray(this.objectRefs);
		ArayHelper.zeroArray(this.candidateRefs);
		ArayHelper.zeroArray(this.candidateDepth);
		ArayHelper.zeroArray(this.children);
		ArayHelper.zeroArray(this.descendants);
		ArayHelper.zeroArray(this.toBottom);
	}

	public void reportBalance() {
		Log.write("TraceBranch report, size=");
		Log.write(this.size);
		Log.write(", set=");
		Log.write(this.setPointers);
		Log.write(", maxDepth=");
		Log.write(this.maxDepth);
		Log.write(", lost=");
		Log.writeln(this.lost);
		Log.flush();
		this.dump();
	}
	
	@Inline
	public void processObject(ObjectReference objRef, long objLevel, int objChildren) {
		if(DEBUG) {
			Log.write("\tTraceBranch.process, objRef=");Log.write(objRef);
    		Log.write(", objLevel=");Log.write(objLevel);
    		Log.write(", objChildren=");Log.write(objChildren);
			Log.write(", size=");Log.writeln(size);
			Log.flush();
		}
		
		int parentIndex = findParent(this.size, objLevel);
		
		if(parentIndex == -1) {
			//root object's descendants tree is exhausted, process it now and remove all slots
			processSubtree(0, this.size - 1);
			if(objChildren > 0) {
				add(objRef, objLevel, objChildren);
			}
			return;
		}
		
		if(objChildren == 0) {
			//we do not keep such nodes in tracker, just update their parent
			addLeafToNode(parentIndex);
			if(!isLast(parentIndex)) {
				processSubtree(parentIndex+1, this.size - 1);
			}
			return;
		}
		
		if(objChildren > Options.sbMaxRefs.getValue()) {
			//we do not keep such nodes in tracker, just update their parent
			addToNode(parentIndex, objChildren);
			return;
		}
		
		if(!isLast(parentIndex)) {
			processSubtree(parentIndex+1, this.size - 1);
		}
		
		add(objRef, objLevel, objChildren);
	}

	/**
	 * Finds the index where the parent (or the closest ancestor) of this node is stored
	 * If the parent is not in the table, this node is not an descendant 
	 * of the currently tracked nodes. In such a case, all the table has to be processed and 
	 * emptied.
	 * 
	 * @param objLevel - current node's level in the global stack
	 * @return parent/ancestor index in table, -1 if not found
	 */
	@Inline
	private int findParent(int index, long objLevel) {
		int result = -1;
		if(DEBUG) {
			Log.write("TraceBranch.findParent, index=");Log.write(index);
			Log.write(", objLevel=");Log.write(objLevel);
		}
		index = index - 1;
		while(index >= 0) {
			if(level[index] <= objLevel) {
				result = index;
				break;
			}
			index--;
		}		
		if(DEBUG) {
			Log.write(", result=");Log.writeln(result);
			Log.flush();
		}
		return result;
	}
	
	@Inline
	private void processSubtree(int from, int to) {	
		if(DEBUG) {
			Log.write("TraceBranch.processSubtree, from=");Log.write(from);
			Log.write(", to=");Log.writeln(to);
			Log.flush();
		}
		
		if(from > to) {
			return;
		}
		//go up accumulating counters in subtree
		int index = to;
		while(index >= from) {
			int parent = findParent(index, this.level[index]);
			if(parent == -1) {
				break;
			}
			addSubtreeToNode(parent, index);
			index--;
		}
		//remember subtree objects
		int distanceParam = Options.sbDistance.getValue();
		long toRemember = (this.toBottom[from] > distanceParam) ? distanceParam : this.toBottom[from];
		for(int i = 0; i < toRemember; i++) {
			rememberObject(from+i);
		}		
		//analyze subtree dimensions, add pointers if needed
		long subtreeSize = this.descendants[from];
		double subtreeRatio = (double)this.toBottom[from]/this.descendants[from];
		if(subtreeSize >= Options.sbSize.getValue() && subtreeRatio > Options.sbRatio.getValue()) {
			if(DEBUG) {
				Log.write("TraceBranch.processSubtree, subtreeSize=");Log.write(subtreeSize);
				Log.write(", subtreeRatio=");Log.writeln(subtreeRatio);
				Log.flush();
			}
		//if(subtreeSize >= PARAMETER_SIZE) {
			addCandidatePointers(from, to);
		}		
		//discard subtree objects
		discardSubtree(from);	
	}
		
	@Inline
	private void rememberObject(int index) {	
		long totalDepth = this.toBottom[index] + index + 1;
		int target = index - Options.sbDistance.getValue();
		if(DEBUG) {
			Log.write("TraceBranch.rememberObject, index=");Log.write(index);
			Log.write(", totalDepth=");Log.write(totalDepth);
			Log.write(", target=");Log.writeln(target);
			Log.flush();
		}
		if(target > 0) {
			if(this.candidateRefs.get(target) == ObjectReference.nullReference() || 
					this.candidateDepth[target] < totalDepth) {
				//replacing candidate ref
				this.candidateRefs.set(target, this.objectRefs.get(index));
				this.candidateDepth[target] = totalDepth;
			} else {
				//forget this object
			}
		}		
	}
	
	@Inline
	private void addCandidatePointers(int from, int to) {
		int distanceParam = Options.sbDistance.getValue();
		int skipParam = Options.sbSkip.getValue();
		
		if(DEBUG) {
			Log.write("TraceBranch.addCandidatePointers, from=");Log.write(from);
			Log.write(" to=");Log.writeln(to);
			Log.flush();
		}
		int index = from;
		//look for the set candidate
		while(index < to) {
			while(this.candidateRefs.get(index) == ObjectReference.nullReference() && index < to) {
				index++;
			}
			if(index < to) {
				ObjectReference obj = this.objectRefs.get(index);
				ObjectReference add = this.candidateRefs.get(index);
				VM.objectModel.setKathy(obj, add);
				this.setPointers++;
			}
			index += skipParam;
		}		
		//add new 
		index -= skipParam;
		int target;
		while((target = index - distanceParam) >= from) {
			if(this.candidateRefs.get(target) == ObjectReference.nullReference()) {
				ObjectReference obj = this.objectRefs.get(target);
				ObjectReference add = this.objectRefs.get(index);
				VM.objectModel.setKathy(obj, add);
				this.setPointers++;
			}
			index -= skipParam;
		}
	}
	
	@Inline
	private void discardSubtree(int from) {
		if(DEBUG && Options.scaleDebug.getValue() > 3) {
			Log.write("TraceBranch.discardSubtree, from=");Log.writeln(from);
			Log.flush();
		}
		if(from == 0) {
			if(DEBUG) {
				Log.write("TraceBranch - discardSubtree, maxDepth=");Log.write(maxDepth);
				Log.write(", this.toBottom[from]=");Log.writeln(this.toBottom[from]);
				Log.flush();
			}
			if(this.toBottom[from] > this.maxDepth) {
				this.maxDepth = this.toBottom[from];
			}
		}
		if(DEBUG) {
			this.reportBalance();
		}
		int index = this.size - 1;
		while(index >= from) {
			discardObject(index);
			index--;
		}
	}
	
	@Inline
	private void discardObject(int index) {
		if(DEBUG) {
			Log.write("TraceBranch.discardObject, index=");
			Log.writeln(index);
			Log.flush();
		}
		this.level[index] = 0;
		this.objectRefs.set(index, ObjectReference.nullReference());
		this.children[index] = 0;
		this.descendants[index] = 0;
		this.toBottom[index] = 0;
		this.size--;
	}
	
	@Inline
	private void addLeafToNode(int node) {
		if(DEBUG) {
			Log.write("TraceBranch.addLeafToNode, node=");Log.writeln(node);
			Log.flush();
		}
		this.descendants[node]++;
		if(this.toBottom[node] == 0) {
			this.toBottom[node] = 1;
		}
	}
	
	@Inline
	private void addToNode(int node, int children) {
		if(DEBUG) {
			Log.write("TraceBranch.addToNode, node=");Log.write(node);
			Log.write(", children=");Log.writeln(children);
			Log.flush();
		}
		this.descendants[node] += children;
		if(this.toBottom[node] == 0) {
			this.toBottom[node] = 1;
		}
	}
	
	@Inline
	private void addSubtreeToNode(int node, int sub) {
		if(DEBUG) {
			Log.write("TraceBranch.addSubtreeToNode, node=");Log.write(node);
			Log.write(" sub=");Log.writeln(sub);
			Log.flush();
		}
		this.descendants[node] += this.descendants[sub] + 1;
		long toBottom = this.toBottom[sub]+1;
		if(this.toBottom[node] < toBottom ) {
			this.toBottom[node] = toBottom;
		}
	}
	
	@Inline
	void dump() {
		Log.writeln("# ObjectRef Level Kids Desc toBottom");
		for (int i = 0; i < this.size; i++) {
			dumpEntry(i);
		}
	}
	
	@Inline
	void dumpEntry(int i) {
		ObjectReference object = objectRefs.get(i);
		Log.write(i);
		Log.write(" ");
		Log.write(object);
		Log.write(" ");
		Log.write(level[i]);
		Log.write(" ");
		Log.write(children[i]);
		Log.write(" ");
		Log.write(descendants[i]);
		Log.write(" ");
		Log.write(toBottom[i]);
		Log.flush();
		VM.objectModel.dumpObject(object);
		Log.flush();
	}

	@Inline
	private boolean isFull() {
		return this.size == HISTORY_SIZE;
	}
	
	@Inline
	private boolean isLast(int index) {
		return this.size == index + 1;
	}

	// push assuming not full
	@Inline
	private void add(ObjectReference objectRef, long level, int children) {
		if(DEBUG) {
			Log.writeln("TraceBranch.add");
			Log.flush();
		}
		if (isFull()) {
			// TODO
			// should not occur when overflow() is implemented
			Log.writeln("TraceBranch.add: can not add object !!!");
			this.lost++;
			//if (VM.VERIFY_ASSERTIONS) {
			//	VM.assertions._assert(!isFull());
			//}
			return;
		}
		
		objectRefs.set(this.size, objectRef);
		this.level[this.size] = level;
		this.children[this.size] = children;
		this.descendants[this.size] = 0;
		this.toBottom[this.size] = 0;
		this.size++;
		if (isFull()) {
			overflow();
		}	
	}

	@Inline
	private void overflow() {
		// TODO need to implement the overflow handler
		Log.writeln("!!!TraceBranch.overflow - do nothing");
		/*
		Log.writeln("!!!TraceBranch.overflow - try to discard");
		Log.flush();
		this.discardSubtree(0);
		Log.writeln("!!!TraceBranch.overflow - done");
		*/
		Log.flush();
	}
}
